#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/1/17 14:09
# @Author  : 杨义乐
# @File    : train.py
# @Description : 训练YOLOv8模型
from ultralytics import YOLO
import os

dataset_path = r'C:\Users\10067\PycharmProjects\pipeline'

# 假设你的训练和验证图像分别存储在 train 和 val 目录下
train_images_path = os.path.join(dataset_path)
val_images_path = os.path.join(dataset_path)

yaml_content = f"""
train: {train_images_path}
val: {val_images_path}
nc: 1
names: ['pip1']
"""

yaml_file_path = os.path.join(dataset_path, 'dataset.yaml')
with open(yaml_file_path, 'w') as f:
    f.write(yaml_content)

# Train YOLOv8
model = YOLO('yolov8n.pt')  # Load a pretrained model (recommended for training)
results = model.train(data=yaml_file_path, epochs=100, imgsz=640, save=True)
